import os
import csv
import shutil
import cv2
import numpy as np

# -------------------- 配置 --------------------
DL_info = 'DL_info.csv'        # DeepLesion 的信息文件
IMG_DIR = 'Images_png'     # 原始图片文件夹
OUT_DIR = 'dataset_yolo'       # 输出 YOLO 数据集根目录
# ---------------------------------------------

# 创建输出文件夹结构
for split in ['images', 'labels']:
    for subset in ['train', 'val', 'test']:
        os.makedirs(os.path.join(OUT_DIR, split, subset), exist_ok=True)

# 读取 CSV
with open(DL_info, 'r', encoding='utf-8') as f:
    reader = csv.reader(f)
    header = next(reader)

    for row in reader:
        file_name = row[0]          # 文件名: '000001_01_01_109.png'
        bbox_str = row[6]           # bbox 坐标: '[x1, y1, x2, y2]'
        split_code = int(row[17])   # 1=train, 2=val, 3=test

        split = {1: 'train', 2: 'val', 3: 'test'}.get(split_code, 'train')

        # 构造路径
        parts = file_name.split('_')
        folder_name = f"{parts[0]}_{parts[1]}_{parts[2]}"
        img_name = parts[3]
        img_path = os.path.join(IMG_DIR, folder_name, img_name)

        if not os.path.exists(img_path):
            continue

        # ---------- 图像读取与处理 ----------
        img = cv2.imread(img_path, -1)
        if img is None:
            continue

        # --- CT HU值预处理 ---
        img = img.astype('float32') - 32768
        img = np.clip(img, -1000, 400)
        img = (img + 1000) / 1400           # 归一化到 [0, 1]
        img = (img * 255).astype('uint8')   # 转为 8-bit
        img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)  # 三通道

        h, w = img.shape[:2]
        # 显示图像
        # ---------- 处理 bbox ----------
        try:
            x1, y1, x2, y2 = map(float, bbox_str.strip('[]').split(','))
        except Exception:
            continue

        x_center = ((x1 + x2) / 2) / w
        y_center = ((y1 + y2) / 2) / h
        bw = (x2 - x1) / w
        bh = (y2 - y1) / h

        # ---------- 写 YOLO 标签 ----------
        label_name = os.path.splitext(file_name)[0] + '.txt'
        label_path = os.path.join(OUT_DIR, 'labels', split, label_name)
        with open(label_path, 'a') as lf:
            lf.write(f"0 {x_center:.6f} {y_center:.6f} {bw:.6f} {bh:.6f}\n")

        # ---------- 保存处理好的图像 ----------
        out_img_path = os.path.join(OUT_DIR, 'images', split, file_name)
        if not os.path.exists(out_img_path):
            cv2.imwrite(out_img_path, img)

print("✅ 单类别 YOLO 病灶检测数据集（已处理图像）生成完成！")
